#include "videoencode.h"

#ifndef DEF_MPP_EN

static VpuCodecContext_t   * ctx = nullptr;
static EncParameter_t      * enc_param = nullptr;
static EncInputStream_t      stream;
static EncInputStream_t    * enc_in = nullptr;
static EncoderOut_t          enc_out_yuv;
static EncoderOut_t        * enc_out = nullptr;

static FILE  * fileOut = nullptr;
static RK_S64 fakeTimeUs = 0;
static RK_S32 size = 0, nal = 0x00000001;

static int width = 0, height = 0, frameRate = 0;
static bool openFlag = false;

namespace VideoEncode
{

bool isOpen()
{
    return (openFlag == true);
}

bool open(const QString &file, int w, int h, int rate)
{
    RK_S32 ret;

    if (file.isNull()) {
        return false;
    }
    width     = w;
    height    = h;
    frameRate = rate;

//    char * mp4file = file.toLocal8Bit().data();
//    ret = (MP4_Write_Init(mp4file, rate) == true) ? 0:-1;
//    if(ret != 0) {
//        qDebug()<<"MP4_Write_Init failed..";
//        return false;
//    }

    fileOut = fopen(file.toStdString().data(), "wb");
    if (fileOut == nullptr){
        return false;
    }
    ret = initMpp();
    if (ret != 0) {
//        MP4_Write_Exit();
        fclose(fileOut);
        fileOut = nullptr;

        return false;
    }

    return openFlag = true;
}

int initMpp()
{
    RK_S32 ret;

    memset(&stream, 0, sizeof(EncInputStream_t));
    enc_in = &stream;
    enc_in->size = 0;
    enc_in->buf = nullptr;

    memset(&enc_out_yuv, 0, sizeof(EncoderOut_t));
    enc_out = &enc_out_yuv;

    ctx = (struct VpuCodecContext *)malloc(sizeof(struct VpuCodecContext));
    if (!ctx){
        return -1;
    }
    memset(ctx, 0, sizeof(struct VpuCodecContext));

    ret = vpu_open_context(&ctx);
    if (ret || ctx == nullptr) {
        return false;
    }

    ctx->codecType      = CODEC_ENCODER;
    ctx->videoCoding    = OMX_RK_VIDEO_CodingAVC;
    ctx->width          = width;
    ctx->height         = height;
    ctx->no_thread      = false;
    ctx->private_data   = malloc(sizeof(EncParameter_t));
    memset(ctx->private_data, 0, sizeof(EncParameter_t));
    ctx->extra_cfg.yuv_format = 0;
    ctx->extra_cfg.bit_depth  = 8;

    enc_param = (EncParameter_t*)ctx->private_data;
    enc_param->width        = width;
    enc_param->height       = height;
    enc_param->format       = ENC_INPUT_YUV420_PLANAR;
    enc_param->rc_mode      = 1;
    enc_param->bitRate      = 4000000; //4M   [10k - 20M]
    enc_param->framerate    = frameRate;
    enc_param->enableCabac  = false;
    enc_param->cabacInitIdc = 0;
    enc_param->intraPicRate = 30;
    enc_param->profileIdc   = 100;
    enc_param->levelIdc     = 40;
    enc_param->qp           = 35;

    if (ctx->init(ctx, nullptr, 0) != 0) {
        return false;
    }
    nal = BSWAP32(nal);
    //写入视频头数据
    if(fileOut &&  ctx->extradata_size > 0) {
        fwrite((uint8_t*)&nal, 1, 4, fileOut);
        fwrite(ctx->extradata, 1, ctx->extradata_size, fileOut);
        //fflush(fileOut);
//        MP4_Set_Nal(nal);
//        MP4_Write_SPS_PPS((uint8_t*)ctx->extradata, ctx->extradata_size);
    }

    ret = ctx->control(ctx, VPU_API_ENC_SETFORMAT, &(enc_param->format));
    if (ret){
        qDebug()<<"VPU_API_ENC_SETFORMAT ret "<<ret;
    }

    ret = ctx->control(ctx, VPU_API_ENC_GETCFG, enc_param);
    if (ret){
        qDebug()<<"VPU_API_ENC_GETCFG ret " << ret;
    }
    enc_param->rc_mode      = 1;
    ret = ctx->control(ctx, VPU_API_ENC_SETCFG, enc_param);
    if (ret){
        qDebug()<<"VPU_API_ENC_SETCFG ret" << ret;
    }

    RK_U32 w_align = ((ctx->width + 15) & (~15));
    RK_U32 h_align = ((ctx->height + 15) & (~15));
    size = w_align * h_align * 3 / 2;

    return 0;
}

bool write(void * buff, int bufSize)
{
    if (buff == nullptr){
        return false;
    }
    //首帧
    enc_in->bufPhyAddr  = (RK_U32)buff;
    enc_in->buf         = (RK_U8 *)buff;
    enc_in->size        = bufSize;
    enc_in->timeUs      = fakeTimeUs;
    fakeTimeUs         += (1000000LL/enc_param->framerate);//40000;

    if (ctx->encoder_sendframe(ctx, enc_in) != 0){
        return false;
    }
    if (ctx->encoder_getstream(ctx, enc_out) != 0){
        return false;
    }
    enc_in->size = 0;
    if ((enc_out->size) && (enc_out->data)) {
        //fwrite((uint8_t*)&nal, 1, 4, pOutFile);
        //MP4_Set_Nal(nal);
        //MP4_Write_Main((uint8_t*)enc_out->data, enc_out->size);
        if(fileOut){
            fwrite((uint8_t*)&nal, 1, 4, fileOut);
            fwrite(enc_out->data, 1, enc_out->size, fileOut);
            //fflush(fileOut);
        }
        enc_out->size = 0;
    }

    return true;
}

void destroyMpp()
{
    ctx->flush(ctx);
    //if (enc_in && enc_in->buf) {
    //    free(enc_in->buf);
    //    enc_in->buf = nullptr;
    //}
    if (enc_out && (enc_out->data)) {
        free(enc_out->data);
        enc_out->data = nullptr;
    }
    if (ctx)
    {
        if (ctx->private_data) {
            free(ctx->private_data);
            ctx->private_data = nullptr;
        }
        vpu_close_context(&ctx);
        ctx = nullptr;
    }
}

void close()
{
    openFlag    = false;
    nal         = 0x00000001;
    fakeTimeUs  = 0;
    destroyMpp();
    if (fileOut) {
        fclose(fileOut);
        fileOut = nullptr;
    }
    //MP4_Write_Exit();
    sync();
}

}

#else
//=====================================================================
//====================================MPP==============================

static struct MPP_ENC_DATA mpp_enc_data;
static int width = 0, height = 0, frameRate = 0;

static bool openFlag = false;

namespace VideoEncode
{

bool open(const QString &file, int w, int h, int rate)
{
    int ret = 0;
    char* mp4file = file.toLocal8Bit().data();

    ret = (MP4_Write_Init(mp4file, rate) == true) ? 0:-1;
    if(ret!=0) {
        qDebug()<<"MP4_Write_Init failed..";
        return false;
    }
    width     = w;
    height    = h;
    frameRate = rate;

    ret = initMpp();
    if (ret != 0) {
        MP4_Write_Exit();
        return false;
    }

    return (openFlag = true);
}

int initMpp()
{
    MPP_RET ret = MPP_OK;
    MppPollType timeout = MPP_POLL_BLOCK;

    memset(&mpp_enc_data, 0, sizeof(mpp_enc_data));
    mpp_enc_data.width      = width;
    mpp_enc_data.height     = height;
    mpp_enc_data.hor_stride = MPP_ALIGN(mpp_enc_data.width, 16);
    mpp_enc_data.ver_stride = MPP_ALIGN(mpp_enc_data.height, 16);
    mpp_enc_data.fmt = MPP_FMT_YUV422_YUYV;
    mpp_enc_data.frame_size = mpp_enc_data.hor_stride * mpp_enc_data.ver_stride * 2;

    mpp_enc_data.type = MPP_VIDEO_CodingAVC;

    ret = mpp_buffer_group_get_internal(&(mpp_enc_data.buf_grp), MPP_BUFFER_TYPE_DRM);
    if (ret){
        qDebug() << "failed to get buffer group internal ret "  <<  ret;
        goto MPP_INIT_OUT;
    }

    ret = mpp_buffer_get(mpp_enc_data.buf_grp, &mpp_enc_data.frm_buf, mpp_enc_data.frame_size);
    if (ret) {
        qDebug() << "failed to get buffer for input frame ret "  <<  ret;
        goto MPP_INIT_OUT;
    }

    ret = mpp_buffer_get(mpp_enc_data.buf_grp, &mpp_enc_data.pkt_buf, mpp_enc_data.frame_size);
    if (ret) {
        qDebug() << "failed to get pkt buffer internal ret "  <<  ret;
        goto MPP_INIT_OUT;
    }
    //创建MPP解码器上下文,初始化mpi
    ret = mpp_create(&mpp_enc_data.ctx, &mpp_enc_data.mpi);
    if (ret) {
        qDebug() << "mpp_create failed ret " << ret;
        goto MPP_INIT_OUT;
    }

    ret = mpp_enc_data.mpi->control(mpp_enc_data.ctx, MPP_SET_OUTPUT_TIMEOUT, &timeout);
    if (MPP_OK != ret) {
        qDebug() << "mpp control timeout ret " << ret;
        goto MPP_INIT_OUT;
    }

    ret = mpp_init(mpp_enc_data.ctx, MPP_CTX_ENC, mpp_enc_data.type);
    if (ret) {
        qDebug() << "mpp_init failed ret " << ret;
        goto MPP_INIT_OUT;
    }
    //帧率
    mpp_enc_data.fps = frameRate;
    mpp_enc_data.gop = 60;
    //像素质量
    mpp_enc_data.bps = (mpp_enc_data.width * mpp_enc_data.height) / 8 * mpp_enc_data.fps;

    mpp_enc_data.prep_cfg.change        = MPP_ENC_PREP_CFG_CHANGE_INPUT | MPP_ENC_PREP_CFG_CHANGE_ROTATION | MPP_ENC_PREP_CFG_CHANGE_FORMAT;
    mpp_enc_data.prep_cfg.width         = mpp_enc_data.width;
    mpp_enc_data.prep_cfg.height        = mpp_enc_data.height;
    mpp_enc_data.prep_cfg.hor_stride    = mpp_enc_data.hor_stride;
    mpp_enc_data.prep_cfg.ver_stride    = mpp_enc_data.ver_stride;
    mpp_enc_data.prep_cfg.format        = mpp_enc_data.fmt;
    mpp_enc_data.prep_cfg.rotation      = MPP_ENC_ROT_0;
    ret = mpp_enc_data.mpi->control(mpp_enc_data.ctx, MPP_ENC_SET_PREP_CFG, &mpp_enc_data.prep_cfg);
    if (ret) {
        qDebug() << "mpi control enc set prep cfg failed ret " << ret;
        goto MPP_INIT_OUT;
    }

    mpp_enc_data.rc_cfg.change  = MPP_ENC_RC_CFG_CHANGE_ALL;
    mpp_enc_data.rc_cfg.rc_mode = MPP_ENC_RC_MODE_CBR;
    mpp_enc_data.rc_cfg.quality = MPP_ENC_RC_QUALITY_MEDIUM;
    mpp_enc_data.rc_cfg.bps_target = mpp_enc_data.bps;

    switch(mpp_enc_data.rc_cfg.rc_mode)
    {
    case MPP_ENC_RC_MODE_FIXQP:
    {
        /* do not setup bitrate on FIXQP mode */
    }break;
    case MPP_ENC_RC_MODE_CBR:
    {
        /* constant bitrate has very small bps range of 1/16 bps */
        mpp_enc_data.rc_cfg.bps_max      = mpp_enc_data.bps * 17 / 16;
        mpp_enc_data.rc_cfg.bps_min      = mpp_enc_data.bps * 15 / 16;
    }break;
    case MPP_ENC_RC_MODE_VBR:
    case MPP_ENC_RC_MODE_AVBR:
    {
        /* VBR mode has wide bound */
        mpp_enc_data.rc_cfg.bps_max      = mpp_enc_data.bps * 17 / 16;
        mpp_enc_data.rc_cfg.bps_min      = mpp_enc_data.bps * 1 / 16;
    }break;
    default:
    {
        /* default use CBR mode */
        mpp_enc_data.rc_cfg.bps_max      = mpp_enc_data.bps * 17 / 16;
        mpp_enc_data.rc_cfg.bps_min      = mpp_enc_data.bps * 15 / 16;
    }break;
    }

    /* fix input / output frame rate */
    mpp_enc_data.rc_cfg.fps_in_flex      = 0;
    mpp_enc_data.rc_cfg.fps_in_num       = mpp_enc_data.fps;
    mpp_enc_data.rc_cfg.fps_in_denorm    = 1;
    mpp_enc_data.rc_cfg.fps_out_flex     = 0;
    mpp_enc_data.rc_cfg.fps_out_num      = mpp_enc_data.fps;
    mpp_enc_data.rc_cfg.fps_out_denorm   = 1;

    mpp_enc_data.rc_cfg.gop              = mpp_enc_data.gop;
    mpp_enc_data.rc_cfg.skip_cnt         = 0;
    mpp_enc_data.rc_cfg.drop_mode        = MPP_ENC_RC_DROP_FRM_DISABLED;
    mpp_enc_data.rc_cfg.drop_threshold   = 20;
    mpp_enc_data.rc_cfg.drop_gap         = 1;

    ret = mpp_enc_data.mpi->control(mpp_enc_data.ctx, MPP_ENC_SET_RC_CFG, &mpp_enc_data.rc_cfg);
    if (ret) {
        qDebug() << "mpi control enc set rc cfg failed ret " << ret;
        goto MPP_INIT_OUT;
    }

    mpp_enc_data.codec_cfg.coding = mpp_enc_data.type;
    switch (mpp_enc_data.codec_cfg.coding)
    {
    case MPP_VIDEO_CodingAVC:
    case MPP_VIDEO_CodingHEVC:
    {
        mpp_enc_data.codec_cfg.h264.change = MPP_ENC_H264_CFG_CHANGE_PROFILE |
                MPP_ENC_H264_CFG_CHANGE_ENTROPY |
                MPP_ENC_H264_CFG_CHANGE_TRANS_8x8;
        /*
         * H.264 profile_idc parameter
         * 66  - Baseline profile
         * 77  - Main profile
         * 100 - High profile
         */
        mpp_enc_data.codec_cfg.h264.profile  = 100;
        /*
         * H.264 level_idc parameter
         * 10 / 11 / 12 / 13    - qcif@15fps / cif@7.5fps / cif@15fps / cif@30fps
         * 20 / 21 / 22         - cif@30fps / half-D1@@25fps / D1@12.5fps
         * 30 / 31 / 32         - D1@25fps / 720p@30fps / 720p@60fps
         * 40 / 41 / 42         - 1080p@30fps / 1080p@30fps / 1080p@60fps
         * 50 / 51 / 52         - 4K@30fps
         */
        if(width==1280 && height==720){
            mpp_enc_data.codec_cfg.h264.level    = 31;
        }else if(width==1920 && height==1080) {
            mpp_enc_data.codec_cfg.h264.level    = 40;//41;
        }

        mpp_enc_data.codec_cfg.h264.entropy_coding_mode  = 1;
        mpp_enc_data.codec_cfg.h264.cabac_init_idc       = 0;
        mpp_enc_data.codec_cfg.h264.transform8x8_mode    = 1;

        mpp_enc_data.codec_cfg.h264.qp_init              = 26;
        mpp_enc_data.codec_cfg.h264.qp_max               = 51;
        mpp_enc_data.codec_cfg.h264.qp_min               = 10;
        mpp_enc_data.codec_cfg.h264.qp_max_i             = 51;//46;
        mpp_enc_data.codec_cfg.h264.qp_min_i             = 10;//24;
        mpp_enc_data.codec_cfg.h264.qp_delta_ip          = 2;
    }break;
    case MPP_VIDEO_CodingMJPEG :
    {
        mpp_enc_data.codec_cfg.jpeg.change  = MPP_ENC_JPEG_CFG_CHANGE_QP;
        mpp_enc_data.codec_cfg.jpeg.quant   = 10;
    } break;
    case MPP_VIDEO_CodingVP8 :
    default :
    {
        qDebug() << "support encoder coding type " << mpp_enc_data.codec_cfg.coding;
    }break;
    }//switch (mpp_enc_data.codec_cfg.coding)

    ret = mpp_enc_data.mpi->control(mpp_enc_data.ctx, MPP_ENC_SET_CODEC_CFG, &mpp_enc_data.codec_cfg);
    if (ret) {
       qDebug() << "mpi control enc set codec cfg failed ret " << ret;
        goto MPP_INIT_OUT;
    }

    /* optional */
    mpp_enc_data.sei_mode = MPP_ENC_SEI_MODE_ONE_FRAME;
    ret = mpp_enc_data.mpi->control(mpp_enc_data.ctx, MPP_ENC_SET_SEI_CFG, &mpp_enc_data.sei_mode);
    if (ret) {
        qDebug() << "mpi control enc set sei cfg failed ret " << ret;
        goto MPP_INIT_OUT;
    }

    if (mpp_enc_data.type == MPP_VIDEO_CodingAVC || mpp_enc_data.type == MPP_VIDEO_CodingHEVC)
    {
        MppPacket packet = nullptr;

        /*
         * Can use packet with normal malloc buffer as input not pkt_buf.
         * Please refer to vpu_api_legacy.cpp for normal buffer case.
         * Using pkt_buf buffer here is just for simplifing demo.
         */
        mpp_packet_init_with_buffer(&packet, mpp_enc_data.pkt_buf);
        /* NOTE: It is important to clear output packet length!! */
        mpp_packet_set_length(packet, 0);

        ret = mpp_enc_data.mpi->control(mpp_enc_data.ctx, MPP_ENC_GET_EXTRA_INFO, &packet);
        if (ret) {
            qDebug() << "mpi control enc get extra info failed";
            goto MPP_INIT_OUT;
        }

        /* get and write sps/pps for H.264 */
        if (packet)
        {
            void *ptr   = mpp_packet_get_pos(packet);
            size_t len  = mpp_packet_get_length(packet);

            MP4_Write_SPS_PPS((uint8_t*)ptr,len);
            //mpp_packet_deinit(&packet);
            packet = nullptr;
        }
    }

    return 0;

MPP_INIT_OUT:
    if (mpp_enc_data.ctx) {
        mpp_destroy(mpp_enc_data.ctx);
        mpp_enc_data.ctx = nullptr;
    }

    if (mpp_enc_data.frm_buf) {
        mpp_buffer_put(mpp_enc_data.frm_buf);
        mpp_enc_data.frm_buf = nullptr;
    }

    if (mpp_enc_data.pkt_buf) {
        mpp_buffer_put(mpp_enc_data.pkt_buf);
        mpp_enc_data.pkt_buf = nullptr;
    }

    if (mpp_enc_data.buf_grp) {
        mpp_buffer_group_put(mpp_enc_data.buf_grp);
        mpp_enc_data.buf_grp = nullptr;
    }

    qDebug() << "init mpp failed!";
    return -1;
}

bool write(void * buff, int bufSize)
{
    MPP_RET ret = MPP_OK;
    MppFrame frame = nullptr;
    MppPacket packet = nullptr;
    MppMeta meta = nullptr;
    void *buf = mpp_buffer_get_ptr(mpp_enc_data.frm_buf);

    //TODO: improve performance here?
    memcpy(buf, buff, bufSize);

    ret = mpp_frame_init(&frame);
    if (ret) {
        qDebug() << "mpp_frame_init failed";
        return -1;
    }

    mpp_frame_set_width(frame, mpp_enc_data.width);
    mpp_frame_set_height(frame, mpp_enc_data.height);
    mpp_frame_set_hor_stride(frame, mpp_enc_data.hor_stride);
    mpp_frame_set_ver_stride(frame, mpp_enc_data.ver_stride);
    mpp_frame_set_fmt(frame, mpp_enc_data.fmt);
    mpp_frame_set_buffer(frame, mpp_enc_data.frm_buf);
    mpp_frame_set_eos(frame, mpp_enc_data.frm_eos);

    meta = mpp_frame_get_meta(frame);
    mpp_packet_init_with_buffer(&packet, mpp_enc_data.pkt_buf);
    /* NOTE: It is important to clear output packet length!! */
    mpp_packet_set_length(packet, 0);
    mpp_meta_set_packet(meta, KEY_OUTPUT_PACKET, packet);

    ret = mpp_enc_data.mpi->encode_put_frame(mpp_enc_data.ctx, frame);
    if (ret)
    {
        qDebug() << "mpp encode put frame failed";
        return -1;
    }
    mpp_frame_deinit(&frame);
    ret = mpp_enc_data.mpi->encode_get_packet(mpp_enc_data.ctx, &packet);
    if (ret)
    {
        qDebug() << "mpp encode get packet failed";
        return -1;
    }
    if(packet)
    {
        // write packet to file here
        void *ptr = mpp_packet_get_pos(packet);
        size_t pkt_len = mpp_packet_get_length(packet);

        mpp_enc_data.pkt_eos = mpp_packet_get_eos(packet);

        //mp4v2 write
        MP4_Write_Main((uint8_t*)ptr,pkt_len);
        mpp_packet_deinit(&packet);

        mpp_enc_data.frame_count++;
        mpp_enc_data.stream_size += pkt_len;

        if (mpp_enc_data.pkt_eos) {
            qDebug() << "found last packet";
        }
    }

/*
    printf("mpp_enc_data.num_frames :%d, frame_count :%d\n",mpp_enc_data.num_frames,mpp_enc_data.frame_count);
    if (mpp_enc_data.num_frames && mpp_enc_data.frame_count >= mpp_enc_data.num_frames) {
        printf("encode max %d frames.\n", mpp_enc_data.frame_count);
        return -1;
    }
*/

    if (mpp_enc_data.frm_eos && mpp_enc_data.pkt_eos) {
        return -1;
    }

    return 0;
}

void destroyMpp()
{
    MPP_RET ret = MPP_OK;

    ret = mpp_enc_data.mpi->reset(mpp_enc_data.ctx);
    if (ret) {
        qDebug() << "mpi->reset failed";
        return ;
    }
    qDebug()<<"destroy mpp";

    if (mpp_enc_data.ctx) {
        mpp_destroy(mpp_enc_data.ctx);
        mpp_enc_data.ctx = nullptr;
    }

    if (mpp_enc_data.frm_buf) {
        mpp_buffer_put(mpp_enc_data.frm_buf);
        mpp_enc_data.frm_buf = nullptr;
    }

    if (mpp_enc_data.pkt_buf) {
        mpp_buffer_put(mpp_enc_data.pkt_buf);
        mpp_enc_data.pkt_buf = nullptr;
    }

    if (mpp_enc_data.buf_grp) {
        mpp_buffer_group_put(mpp_enc_data.buf_grp);
        mpp_enc_data.buf_grp = nullptr;
    }
}

void close()
{
    openFlag = false;
    destroyMpp();
    MP4_Write_Exit();
    sync();
}

}
#endif
